package com.sm2cppa.utils;

import org.bouncycastle.asn1.x9.ECNamedCurveTable;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.math.ec.FixedPointCombMultiplier;
import org.bouncycastle.util.Arrays;
import java.math.BigInteger;
import java.nio.ByteBuffer;


public class KeyDerive {
    private ECDomainParameters ecParams; // SM2曲线参数

    /**
     * 公钥派生算法
     * @param publicKey  SM2公钥
     * @param keyIndex  派生密钥索引
     * @param chain  链码
     * @return 派生的公钥和新的链码
     */
    public static byte[][] SM2_PublicKeyDerive(byte[] publicKey,int keyIndex,byte[] chain) {
        //产生密钥

        X9ECParameters x9 = ECNamedCurveTable.getByName("sm2p256v1");
        BigInteger curveOrder = x9.getN();

        //使用SM3哈希
        byte[]hash=Hashpkichain(publicKey,keyIndex,chain);

        //截取哈希的前半部分和后半部分
        int midindex = hash.length / 2;
        byte[] firstHalf = Arrays.copyOfRange(hash, 0, midindex); // 前半部分
        byte[] secondHalf = Arrays.copyOfRange(hash, midindex, hash.length); // 后半部分
        BigInteger HL = new BigInteger(1, firstHalf);


        //将公钥字节流解码为椭圆曲线点
        ECPoint P= KeyDerive.PktoECpoint(publicKey);
        ECPoint newP = (new FixedPointCombMultiplier()).multiply(P, HL);
        byte[] newPk = newP.getEncoded(false);
        if (newPk.length == 65) {
                newPk = Arrays.copyOfRange(newPk, 1, newPk.length);
        }

        byte[] newChain=secondHalf;
        byte[][] pkandchain=new byte[][]{newPk, newChain};
        return pkandchain;
    }

    /**
     * 私钥派生算法
     * @param privateKey  SM2私钥
     * @param keyIndex  派生密钥索引
     * @param chain  链码
     * @return 派生的公钥和新的链码
     */
    public  static byte[][] SM2_PrivateKeyDerive(byte[] privateKey,int keyIndex,byte[] chain){
        X9ECParameters x9 = ECNamedCurveTable.getByName("sm2p256v1");
        BigInteger curveOrder = x9.getN();
        BigInteger d = new BigInteger(1, privateKey);

        byte[] SM2_PublicKey=SM2_Getpkfromsk(privateKey);

        byte[] hash=Hashpkichain(SM2_PublicKey,keyIndex,chain);
        int midIndex = hash.length / 2;
        byte[] firstHalf = Arrays.copyOfRange(hash, 0, midIndex); // 前半部分
        byte[] secondHalf = Arrays.copyOfRange(hash, midIndex, hash.length); // 后半部分

        BigInteger HL = new BigInteger(1, firstHalf);
        BigInteger skH= d.multiply(HL).mod(curveOrder);//模模数
        byte[] newSk= skH.toByteArray();
        if (newSk.length == 33) {
            newSk = Arrays.copyOfRange(newSk, 1, newSk.length);
        }

        byte[] newChain=secondHalf;
        byte[][] skandchain=new byte[][]{newSk, newChain};
        return skandchain;
    }


    /**
     * 将输入的SM2公钥和派生密钥索引和链码合并后放入SM3哈希函数输出
     * @param SM2_PublicKey  SM2公钥
     * @param keyIndex  派生密钥索引
     * @param chain  链码
     * @return SM3哈希结果
     */
    public static  byte[] Hashpkichain(byte[] SM2_PublicKey,int keyIndex ,byte[] chain) {
        //切割公钥
        int midkeyIndex1 = SM2_PublicKey.length / 2;
        byte[] firstHalfSM2_PublicKey = Arrays.copyOfRange(SM2_PublicKey, 0, midkeyIndex1); // 前半部分
        // 创建哈希实例
        SM3Digest digest = new SM3Digest();
        // 准备输入字节数组
        byte[] iBytes = ByteBuffer.allocate(4).putInt(keyIndex).array();
        byte[] input = new byte[SM2_PublicKey.length + iBytes.length + chain.length];
        // 将公钥、索引i和chainroot合并到输入数组中
        System.arraycopy(SM2_PublicKey, 0, input, 0, SM2_PublicKey.length); // 公钥
        System.arraycopy(iBytes, 0, input, SM2_PublicKey.length, iBytes.length); // 索引i
        System.arraycopy(chain, 0, input, SM2_PublicKey.length + iBytes.length, chain.length); // chainroot
        // 更新哈希
        digest.update(input, 0, input.length);// 32 字节（256 位）
        byte[] hash = new byte[digest.getDigestSize()];
        digest.doFinal(hash, 0);
        return hash;
    }


    /**
     * 使用SM3哈希函数计算公钥哈希
     * @param SM2_PublicKey  SM2公钥
     * @return 公钥哈希
     */
    public static byte[] Sm3Hash(byte[] SM2_PublicKey){
        SM3Digest digest = new SM3Digest();
        digest.update(SM2_PublicKey, 0, SM2_PublicKey.length);// 32 字节（256 位）
        byte[] hash = new byte[digest.getDigestSize()];
        digest.doFinal(hash, 0);
        // System.out.println("Hash (Hex): " + Util.byte2HexStr(hash));
        return hash;
    }


    /**
     * 将公钥字节流转换为椭圆曲线点
     * @param SM2_PublicKey  SM2公钥
     * @return 公钥对应的椭圆曲线点
     */
    public static ECPoint PktoECpoint(byte[] SM2_PublicKey) {
        X9ECParameters x9 = ECNamedCurveTable.getByName("sm2p256v1");
        ECCurve mCurve = x9.getCurve();
        // 添加前缀
        byte[] fullPublicKey = new byte[SM2_PublicKey.length + 1];
        fullPublicKey[0] = 0x04;  // 添加未压缩前缀
        System.arraycopy(SM2_PublicKey, 0, fullPublicKey, 1, SM2_PublicKey.length);
        // 解码公钥点
        return mCurve.decodePoint(fullPublicKey);
    }

    /**
     * 根据输入的私钥计算公钥
     * @param SM2_privateKey SM2私钥
     * @return SM2公钥
     */
    public static byte[] SM2_Getpkfromsk(byte[] SM2_privateKey){
        X9ECParameters x9 = ECNamedCurveTable.getByName("sm2p256v1");
        BigInteger curveOrder = x9.getN();
        BigInteger d = new BigInteger(1, SM2_privateKey);
        ECPoint Gsk = (new FixedPointCombMultiplier()).multiply(x9.getG(),d);
        byte[] SM2_PublicKey = Gsk.getEncoded(false);
        if (SM2_PublicKey.length == 65) {
            SM2_PublicKey = Arrays.copyOfRange(SM2_PublicKey, 1, SM2_PublicKey.length);
        }
        return SM2_PublicKey;
    }

}